#质心计算过程
import numpy as np
import matplotlib.pyplot as plt

plt.figure(figsize=[12, 6])
spr = 2
spc = 3
spn = 0

# 样本集
X = np.array([[1, 2], [2, 2], [6, 8], [7, 8]], dtype=np.float64)  # m x n
#定义初始化质心
C = np.array([[1, 2], [2, 2]], dtype=np.float64)  # n_cluster x n
m = len(X)
n_cluster = len(C)

spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(X[:, 0], X[:, 1], color='b')
plt.scatter(C[:, 0], C[:, 1], color='r', marker='x', s=100, zorder=100)


def x_visualize(X, C, class_vector, title):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    classes = np.unique(class_vector)
    n_cluster = len(classes)
    cmap = plt.cm.get_cmap('rainbow', n_cluster)
    for i, cls in enumerate(classes):
        c = cmap(i)
        idx = class_vector == cls
        plt.scatter(X[idx, 0], X[idx, 1], color=c)
        plt.scatter(C[i, 0], C[i, 1], color=c, marker='x', s=100, zorder=100)


iters = 0
while (iters < 2):
    iters += 1

    # cluster
    dis_mat = np.zeros([m, n_cluster])
    for i in range(n_cluster):
        dis_mat[:, i] = np.sqrt(((X - C[i])**2).sum(axis=1))
    class_vector = dis_mat.argmin(axis=1)

    # visualize
    x_visualize(X, C, class_vector, f'clustering #{iters}')

    # new center
    for i in range(n_cluster):
        C[i] = X[class_vector == i].mean(axis=0)

    # visualize
    x_visualize(X, C, class_vector, f'renew centers #{iters}')

plt.show()
